#include "klee/Support/WICFG.h"
#include <fstream>
// #include <cassert>
#include <sstream>
#include <algorithm>
#include "klee/Support/ErrorHandling.h"
#include <unistd.h>
using namespace llvm;

namespace klee {
CallChain_Node::ID_t CallChain_Node::_id = 0;
ConstraintTreeNode::ID_t ConstraintTreeNode::_id = 0;

void ConstraintTree::removeChild(ConstraintTreeNode::NodePtr& child) {
    std::cout << "removeChild of constraintTree" << std::endl;
    auto parent = child->parent;
    while(parent != nullptr && child->next.empty()) { // 直到整棵树都删光了，或子节点不为叶节点时结束循环
        // parent->next.erase(child); // 先把子节点从父节点的后继节点列表中删除出
        for(auto n = parent->next.begin(); n != parent->next.end(); ++n) {
            if(n->get()->ID == child->ID) {
                parent->next.erase(n);
                break;
            }
        }
        // std::remove_if(parent->next.begin(), parent->next.end(), [&](ConstraintTreeNode::NodePtr it){return child.get() == it.get();}); // 先把子节点从父节点的后继节点列表中删除出
        child = parent;
        parent = parent->parent;
    }
}

void WICFG::init(llvm::Function *targetFunc) {
    // entryNode = make_shared<WICFG_Node>(new WICFG_Node(targetFunc));
    std::cout << "build WICFG" << std::endl;
    // return ;
    entryNode = std::make_shared<WICFG_Node>(targetFunc, nullptr, numOfWICFG_Node++, 0);
    callChain.entryNode = std::make_shared<CallChain_Node>(nullptr);
    endNode = traverseFunc(targetFunc, nullptr, entryNode, callChain.entryNode, true, 0);
    // llvm::raw_os_ostream O(std::cout);
    // targetFunc->print(O);
    std::cout << "Traverse done" << std::endl;
    for(auto n : uncoveredTargetWICFG_Node) {
        weight_t w = initWeight * n->numOfInst;
        std::set<NodePtr> alreadyChangedNode;
        if(n->opcodeNameSeries[n->opcodeNameSeries.size() - 1] == "ret") {
            alreadyChangedNode.insert(n);
            setDistance(n, 0, alreadyChangedNode);
            alreadyChangedNode.clear();
        }
        alreadyChangedNode.insert(n);
        changeWeight(n, w, 0, true, alreadyChangedNode);
        // alreadyChangedNode.clear();
        // alreadyChangedNode.insert(n);
        // changeWeight2(n, w, 0, true, alreadyChangedNode);
    }
    printToDotFile();
    callChain.printToDotFile();
    // for(;;);
}

void WICFG::changeWeight(NodePtr targetNode, weight_t weight, int pathLen, bool isAdd, std::set<NodePtr>& alreadyChangedNode) {
    // if(weight <= 0) { // 权值的绝对值小于0，就不继续递归了
    //     return;
    // }
    targetNode->weight += (isAdd ? weight : -weight);
    for(auto& n : targetNode->predecessors) {
        auto it = alreadyChangedNode.find(n);
        if(it == alreadyChangedNode.end()) {
            alreadyChangedNode.insert(n);
        } else {
            // klee_warning("重复加权");            
            continue;
        }
        // weight_t w = weight / n->successors.size() - pathLen * 5 - n->numOfInst * 2;
        weight_t w = weight - pathLen * 2;
        w = std::max(w, weight_t(1));
        // n->weight += (isAdd ? w : -w);
        changeWeight(n, w, pathLen + 1, isAdd, alreadyChangedNode);
    }
}

void WICFG::setDistance(NodePtr targetNode, int pathLen, std::set<NodePtr>& alreadyChangedNode) {
    targetNode->distanceToEnd = pathLen;
    for(auto& n : targetNode->predecessors) {
        auto it = alreadyChangedNode.find(n);
        if(it == alreadyChangedNode.end()) {
            alreadyChangedNode.insert(n);
        } else {
            continue;
        }
        setDistance(n, pathLen + 1, alreadyChangedNode);
    }
}

// void WICFG::changeWeight2(NodePtr targetNode, weight_t weight, int pathLen, bool isAdd, std::set<NodePtr>& alreadyChangedNode) {
//     // if(weight <= 0) { // 权值的绝对值小于0，就不继续递归了
//     //     return;
//     // }
//     targetNode->weight += (isAdd ? weight : -weight);
//     for(auto& n : targetNode->successors) {
//         auto it = alreadyChangedNode.find(n);
//         if(it == alreadyChangedNode.end()) {
//             alreadyChangedNode.insert(n);
//         } else {
//             // klee_warning("重复加权");            
//             continue;
//         }
//         // weight_t w = weight / n->successors.size() - pathLen * 5 - n->numOfInst * 2;
//         weight_t w = weight;
//         w = std::max(w, weight_t(1));
//         // n->weight += (isAdd ? w : -w);
//         changeWeight2(n, w, pathLen + 1, isAdd, alreadyChangedNode);
//     }
// }

bool WICFG::checkIsEncounterLoop(NodePtr curNode, NodePtr startNode, std::set<NodePtr>& inProcess) { // 递归回溯检测循环
    for(auto n : curNode->successors) {
        if(inProcess.find(n) == inProcess.end()) {
            inProcess.emplace(n);
            if(checkIsEncounterLoop(n, startNode, inProcess)) {
                return true;
            }
            inProcess.erase(n);
        } else if(n == startNode) {
            return true;
        }
    }
    return false;
}

bool WICFG::checkIsEncounterRecursion(CallChain_Node::NodePtr callChainNode, llvm::Function *callee) {
    CallChain_Node::NodePtr predecessorCallChainNode = callChainNode;
    while(predecessorCallChainNode != nullptr) {
        if(callee == predecessorCallChainNode->CallFunc) {
            return true;
        } else {
            predecessorCallChainNode = predecessorCallChainNode->predecessors;
        }
    }
    return false;
}

void WICFG::fillInst2WICFGNodeVec(Inst2WICFGNodeVec_t &mp, llvm::Instruction *BB, NodePtr curNode) {
    // 把新增的WICFG节点关联到对应的指令
    auto it = mp.find(BB);
    if(it != mp.end()) {
        // std::cout << "===================>gettwo" << std::endl;
        mp[BB].emplace_back(curNode);
    } else {
        // std::cout << "===================>getone" << std::endl;
        // std::vector<WICFG_Node*> WNVec(curNode);
        std::vector<NodePtr> WNVec;
        WNVec.emplace_back(curNode);
        mp[BB] = WNVec;
    }
}

void WICFG::fillBB2WICFGNodeVec(BB2WICFGNodeVec_t &mp, llvm::BasicBlock *BB, NodePtr curNode) {
    // 把新增的WICFG节点关联到对应的基本块
    auto it = mp.find(BB);
    if(it != mp.end()) {
        // std::cout << "===================>gettwo" << std::endl;
        mp[BB].emplace_back(curNode);
    } else {
        // std::cout << "===================>getone" << std::endl;
        // std::vector<WICFG_Node*> WNVec(curNode);
        std::vector<NodePtr> WNVec;
        WNVec.emplace_back(curNode);
        mp[BB] = WNVec;
    }
}

WICFG::NodePtr WICFG::traverseBasicBlock(llvm::BasicBlock *BB, llvm::Instruction *callerInst, NodePtr preNode,
            NodePtr &retNode, CallChain_Node::NodePtr callChainNode, bool isTargetFunc, FuncSeries_t funcSeries, bool fromBr) {
    // std::cout << "TraverseBB" << std::endl;
    // cout << preNode->ID << endl;
    if(BB2WICFGNodeMapVec.size() > funcSeries) { // 检查该基本块是否已遍历
        auto node = BB2WICFGNodeMapVec[funcSeries].find(BB);
        if(node != BB2WICFGNodeMapVec[funcSeries].end()) {
            if(fromBr) { // 貌似没起作用？
                preNode->successors.emplace_back(node->second[0]);
                node->second[0]->predecessors.emplace_back(preNode);
            }
            return node->second[0];
        }
    }
    // 创建新WICFG节点，并维护父子节点的链接
    NodePtr curNode = std::make_shared<WICFG_Node>(BB->getParent(), callerInst, numOfWICFG_Node++, funcSeries, isTargetFunc);
    // preNode->succssor.emplace_back(make_shared<WICFG_Node>(curNode));
    // curNode->predecessor.emplace_back(make_shared<WICFG_Node>(preNode));
    preNode->successors.emplace_back(curNode);
    curNode->predecessors.emplace_back(preNode);
    if(isTargetFunc) {
        uncoveredTargetWICFG_Node.insert(curNode);
    }
    // std::cout << "未覆盖目标BB: " << uncoveredTargetWICFG_Node.size() << std::endl;
    // 记录已遍历的基本块对应的WICFG节点，局部于CFG
    if(BB2WICFGNodeMapVec.size() <= funcSeries) {
        // BB2WICFGNode_t map(BB, curNode);
        BB2WICFGNodeVec_t map;
        fillBB2WICFGNodeVec(map, BB, curNode);
        BB2WICFGNodeMapVec.emplace_back(map);
    } else {
        fillBB2WICFGNodeVec(BB2WICFGNodeMapVec[funcSeries], BB, curNode);
    }

    // 把新增的WICFG节点关联到对应的基本块，全局ICFG
    fillBB2WICFGNodeVec(BB2WICFGNodeVecMap, BB, curNode);

    // 遍历当前基本块中的指令
    for(const auto &inst : *BB) {
        fillInst2WICFGNodeVec(Inst2WICFGNodeVecMap, const_cast<llvm::Instruction *>(&inst), curNode);
        // cout << inst.getOpcodeName() << endl;
        curNode->opcodeNameSeries.emplace_back(inst.getOpcodeName());
        curNode->numOfInst++;
        if(isTargetFunc) {
            numOfTargetFuncInst++;
        }
        // inst.print(outs());
        switch(inst.getOpcode()) {
        case Instruction::Call: {
#if LLVM_VERSION_CODE >= LLVM_VERSION(8, 0)
            const CallBase &cs = cast<CallBase>(inst);
#else
            const CallSite cs(inst);
#endif
            Value *calledVal = cs.getCalledValue();
            // unsigned numArgs = cs.arg_size();
            // Function *callee = getTargetFunction(fp, state);
            Constant *c = dyn_cast<Constant>(calledVal);
            GlobalValue *gv = dyn_cast<GlobalValue>(c);
            Function *callee = dyn_cast<Function>(gv);

            // FixMe：后续通过关联callgraph节点来检测循环调用
            const Function *fp = inst.getParent()->getParent();
            assert(fp == callChainNode->CallFunc && "Call Chain Node link with wrong function!");
            if(checkIsEncounterRecursion(callChainNode, callee)) {
                // 递归调用处理，权值计算时需要降权处理。
                curNode->isEncounterRecursion = true;
            } else {
                // std::cout << "Instruction::call" << std::endl;
                preNode = traverseFunc(callee, const_cast<llvm::Instruction*>(&inst), curNode, callChainNode, false, ++numOfCalledTime);
                curNode = std::make_shared<WICFG_Node>(BB->getParent(), callerInst, numOfWICFG_Node++, funcSeries, isTargetFunc);
                curNode->predecessors.emplace_back(preNode);
                preNode->successors.emplace_back(curNode);
                if(isTargetFunc) {
                    uncoveredTargetWICFG_Node.insert(curNode);
                }
                // 记录已遍历的基本块对应的WICFG节点
                // if(BB2WICFGNodeMapVec.size() <= funcSeries) {
                //     // BB2WICFGNode_t map(BB, curNode);
                //     BB2WICFGNodeVec_t map;
                //     map.emplace(BB, {curNode});
                //     BB2WICFGNodeMapVec.emplace_back(map);
                // } else {
                //     auto it = BB2WICFGNodeMapVec[funcSeries].find(BB);
                //     if(it == BB2WICFGNodeMapVec[funcSeries].end()) {
                //         BB2WICFGNodeMapVec[funcSeries][BB] = curNode;
                //     } else {
                //         BB2WICFGNodeMapVec[funcSeries][BB].emplace_back(curNode);
                //     }
                // }
                fillBB2WICFGNodeVec(BB2WICFGNodeMapVec[funcSeries], BB, curNode);

                fillBB2WICFGNodeVec(BB2WICFGNodeVecMap, BB, curNode);
            }
            break;
        }
        case Instruction::Br: {
            // std::cout << "Instruction::Br" << std::endl;
            const BranchInst *bi = cast<BranchInst>(&inst);
            unsigned numOfSuccessors = bi->getNumSuccessors();
            // NodePtr parentNode = curNode;
            // std::vector<NodePtr> childNode;
            for(unsigned i = 0; i < numOfSuccessors; ++i) {
                // BasicBlock *getSuccessor(unsigned i) const;
                llvm::BasicBlock *nextBB = bi->getSuccessor(i);
                preNode = traverseBasicBlock(nextBB, callerInst, curNode, retNode, callChainNode, isTargetFunc, funcSeries, true);
            }
            std::set<NodePtr> inProcess{curNode};
            if(numOfSuccessors > 1 && checkIsEncounterLoop(curNode, curNode, inProcess)) {
                curNode->isEncounterLoop = true;
            }
            break;
        }
        case Instruction::IndirectBr: {
            std::cout << "catch indirectbr" << std::endl;
            break;
        }
        case Instruction::Switch: {
            const SwitchInst *si = cast<SwitchInst>(&inst);      // Iterate through all non-default cases and order them by expressions
            unsigned numOfSuccessors = si->getNumSuccessors();
            for(unsigned i = 0; i < numOfSuccessors; ++i) {
                llvm::BasicBlock *nextBB = si->getSuccessor(i);
                preNode = traverseBasicBlock(nextBB, callerInst, curNode, retNode, callChainNode, isTargetFunc, funcSeries, true);
            }
            break;
        }
        case Instruction::Ret: {
            retNode = curNode;
            break;
        }
        default: {
            // std::cout << "Other Instrution" << std::endl;
            break;
        }
        }
    }
    return curNode;
}

// 遍历函数，返回遍历函数的最后一个WICFG节点
WICFG::NodePtr WICFG::traverseFunc(llvm::Function *fp, llvm::Instruction *callerInst, NodePtr preNode,
                                CallChain_Node::NodePtr preCallChainNode, bool isTargetFunc, FuncSeries_t funcSeries) {
    // std::cout << "TraverseFunc" << std::endl;
    CallChain_Node::NodePtr callChainNode = std::make_shared<CallChain_Node>(fp);
    callChainNode->predecessors = preCallChainNode;
    preCallChainNode->successors.emplace_back(callChainNode);
    NodePtr retNode;
    for(auto &BB : *fp) { // llvm::Function必须是指针然后加星号才能这样遍历，why？？？
        // std::cout << "enter" << std::endl;
        preNode = traverseBasicBlock(&BB, callerInst, preNode, retNode, callChainNode, isTargetFunc, funcSeries, false);
    }
    return retNode;
}

static std::string getIndent() {
    return "    ";
}

std::string WICFG_Node::getCircleBody() {
    std::string head = getIndent() + 
        std::to_string(ID) +
        (isTarget ? " [shape=circle,label=" : " [shape=circle,style=dashed,label=") +
        "\"" + std::to_string(ID) +
        // "\"Node_ID: " + std::to_string(ID) +
        // "\\lFunc: " + func->getName().str() +
        // "\\lCall_Seq: " + std::to_string(callSeries) +
        // "\\lWeight: " + std::to_string(weight) +
        // "\\lInst_Num: " + std::to_string(numOfInst) + 
        "\\l";
    if(this->predecessors.size() == 0) {
        return head + "\"];";
        return head + "Entry\\l\"];";
    }
    std::string body = head;
    // for(auto &opName : opcodeNameSeries) {
    //     body += opName + "\\l";
    // }
    // body += "Last_Inst: " + opcodeNameSeries[opcodeNameSeries.size() - 1] + "\\l";
    if(isEncounterRecursion) {
        body += "Recursion!\\l";
    }
    if(isEncounterLoop) {
        body += "Loop!\\l";
    }
    body += "\"];";
    return body;
}

std::string WICFG_Node::getBody() {
    std::string head = getIndent() + 
        std::to_string(ID) + 
        (isTarget ? " [shape=record,label=" : " [shape=record,style=dashed,label=") +
        "\"{Node_ID: " + std::to_string(ID) +
        "\\lFunc: " + func->getName().str() +
        "\\lCall_Seq: " + std::to_string(callSeries) +
        "\\lWeight: " + std::to_string(weight) +
        "\\lInst_Num: " + std::to_string(numOfInst) + "\\l";
    if(this->predecessors.size() == 0) {
        return head + "EntryNode\\l}\"];";
    }
    std::string body = head;
    // for(auto &opName : opcodeNameSeries) {
    //     body += opName + "\\l";
    // }
    body += "Last_Inst: " + opcodeNameSeries[opcodeNameSeries.size() - 1] + "\\l";
    if(isEncounterRecursion) {
        body += "encounterRecursion!\\l";
    }
    if(isEncounterLoop) {
        body += "encounterLoop!\\l";
    }
    body += "}\"];";
    return body;
}

std::string WICFG::printSuccessor(NodePtr parent) {
    alreadyPrinted.insert(parent);
    std::string body;
    // body += parent->getBody() + "\n";
    body += parent->getCircleBody() + "\n";
    // std::cout << parent->successors.size() << std::endl;
    for(auto &successor : parent->successors) {
        body += getIndent() + std::to_string(parent->ID) + " -> " + std::to_string(successor->ID) + ";\n";
        if(alreadyPrinted.find(successor) == alreadyPrinted.end()) {
            body += printSuccessor(successor);
        }
    }
    return body;
}

std::string WICFG::printPredecessor(NodePtr child) {
    alreadyPrinted.insert(child);
    std::string body;
    for(auto &predecessor : child->predecessors) {
        body += getIndent() + std::to_string(child->ID) + " -> " + std::to_string(predecessor->ID) +";\n";
        if(alreadyPrinted.find(predecessor) == alreadyPrinted.end()) {
            body += printPredecessor(predecessor);
        }
    }
    return body;
}

void WICFG::printToDotFile() {
    static int id = 0;
    llvm::raw_os_ostream O(std::cout);
    std::string head = "digraph \"WICFG\" {\n";
    std::string tail = "}\n";
    std::string body = getIndent() + "label=\"WICFG\";\n";
    body += "ratio=2;\n";
    body += printSuccessor(entryNode);
    alreadyPrinted.clear();
    // body += printPredecessor(endNode);
    // alreadyPrinted.clear();
    std::string out = head + body + tail;
    // std::cout << out << std::endl;
    std::string fileName = "a_" + std::to_string(id++) + ".dot";
    std::fstream fs;
    // fs.open("a.dot", std::ios::out);
    fs.open(fileName.c_str(), std::ios::out);
    fs << out;
    std::cout << "Write WICFG Graph to " << fileName << std::endl;
}

std::string CallChain_Node::getCircleBody() {
    std::string head = getIndent() + 
        std::to_string(ID) +
        (true ? " [shape=circle,label=" : " [shape=circle,style=dashed,label=") +
        "\"" + std::to_string(ID) +
        // "\"Node_ID: " + std::to_string(ID) +
        // "\\lFunc: " + func->getName().str() +
        // "\\lCall_Seq: " + std::to_string(callSeries) +
        // "\\lWeight: " + std::to_string(weight) +
        // "\\lInst_Num: " + std::to_string(numOfInst) + 
        "\\l";
    if(this->predecessors == nullptr) {
        // return head + "\"];";
        return head + "Entry\\l\"];";
    }
    std::string body = head;
    body += "\"];";
    return body;
}

std::string CallChain::printSuccessor(NodePtr parent) {
    alreadyPrinted.insert(parent);
    std::string body;
    // body += parent->getBody() + "\n";
    body += parent->getCircleBody() + "\n";
    for(auto &successor : parent->successors) {
        body += getIndent() + std::to_string(parent->ID) + " -> " + std::to_string(successor->ID) + ";\n";
        if(alreadyPrinted.find(successor) == alreadyPrinted.end()) {
            body += printSuccessor(successor);
        }
    }
    return body;
}

void CallChain::printToDotFile() {
    llvm::raw_os_ostream O(std::cout);
    std::string head = "digraph \"CallGraph\" {\n";
    std::string tail = "}\n";
    std::string body = getIndent() + "label=\"CallChain\";\n";
    // body += "ratio=2;\n";
    body += printSuccessor(entryNode);
    alreadyPrinted.clear();
    std::string out = head + body + tail;
    std::string fileName = "call.dot";
    std::fstream fs;
    fs.open(fileName.c_str(), std::ios::out);
    fs << out;
    std::cout << "Write Call Graph to " << fileName << std::endl;
}

std::string ConstraintTreeNode::getBody() {
    std::string head = getIndent() + 
        std::to_string(ID) + " [shape=circle,label=" +
        "\"" + std::to_string(ID) +
        "\\l";
    if(this->parent == nullptr) {
        return head + "Entry\\l\"];";
    }
    std::string body = head;
    body += "\"];";
    return body;
}

std::string ConstraintTree::printSuccessor(ConstraintTreeNode::NodePtr parent) {
    assert(parent != nullptr);
    std::string body;
    body += parent->getBody();
    for(auto child : parent->next) {
        body += getIndent() + std::to_string(parent->ID) + " -> " + std::to_string(child->ID) + ";\n";
        body += printSuccessor(child);
    }
    return body;
}

void ConstraintTree::printToDotfile() {
    assert(root != nullptr);
    llvm::raw_os_ostream O(std::cout);
    std::string head = "digraph \"CallGraph\" {\n";
    std::string tail = "}\n";
    std::string body = getIndent() + "label=\"ConstraintTree\";\n";
    body += printSuccessor(root);
    std::string out = head + body + tail;
    std::string fileName = "constraintTree.dot";
    std::fstream fs;
    fs.open(fileName.c_str(), std::ios::out);
    fs << out;
    std::cout << "Write Constraint Tree to " << fileName << std::endl;
}
};